package com.sliding.aspectj;

import com.common.utils.IpUtils;
import com.common.utils.ServletUtils;
import com.sliding.annotation.RateLimit;
import com.sliding.enums.ModeType;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.UUID;
import java.util.concurrent.TimeUnit;

@Aspect //告诉框架我是一个切面类
@Component// 随着框架的启动而启动
@Slf4j
@AllArgsConstructor
public class RateLimitAspect {

    private final RedisTemplate<Object, Object> redisTemplate;

    @Around("@annotation(rateLimit)")
    public Object around(ProceedingJoinPoint point, RateLimit rateLimit) throws Throwable {
        // 组合限流 key
        String limitKey = generateKey(point, rateLimit);
        // 执行限流逻辑
        boolean limited = isLimited(limitKey, rateLimit.window(), rateLimit.threshold());
        if (limited) {
            return rateLimit.message();
        }
        // 执行目标方法
        return point.proceed();
    }


    private String generateKey(ProceedingJoinPoint point, RateLimit rateLimit) {
        StringBuilder sb = new StringBuilder(rateLimit.prefix());
        // 获取请求的类名
        String className = point.getTarget().getClass().getName();
        // 获取请求的方法名
        String methodName = point.getSignature().getName();
        if (rateLimit.mode() == ModeType.IP) {
            sb.append(IpUtils.getIpAddress(ServletUtils.getRequest())).append("-");
        }
        sb.append(className).append("-").append(methodName);
        return sb.toString();
    }


    private boolean isLimited(String key, int window, int threshold) {
        long current = System.currentTimeMillis();
        // 删除旧数据，即 window 秒之前的数据
        redisTemplate.opsForZSet().removeRangeByScore(key, 0, current - window * 1000L);

        // 获取当前计数
        Long count = redisTemplate.opsForZSet().zCard(key);

        if (count != null && count >= threshold) {
            return true;
        }

        // 添加请求
        redisTemplate.opsForZSet().add(key, UUID.randomUUID().toString(), current);
        redisTemplate.expire(key, window * 1000L, TimeUnit.SECONDS);
        log.info("限流 key: {}, count: {}, threshold: {}, window: {}", key, count + 1, threshold, window);
        return false;
    }
}
